import torch 
from transformer_model import BertConfig, CLS

import torchvision
import torch
import torch.nn as nn 
from torchvision import datasets, transforms
import matplotlib.pyplot as plt 
from torch.utils.tensorboard import SummaryWriter       

if __name__ == "__main__":

    writer = SummaryWriter("./mnist_log/")

    config = BertConfig(img_size=(28,28), num_hidden_layers=1)
    model = CLS(config, out_size=10)

    transform = transforms.Compose([transforms.ToTensor(),
                               transforms.Normalize(mean=[0.5],std=[0.5])])
    data_train = datasets.MNIST(root = "./data/",
                            transform=transform,
                            train = True,
                            download = True)

    data_test = datasets.MNIST(root="./data/",
                           transform = transform,
                           train = False)

    data_loader_train = torch.utils.data.DataLoader(dataset=data_train,
                                                batch_size = 64,
                                                shuffle = True)

    data_loader_test = torch.utils.data.DataLoader(dataset=data_test,
                                                batch_size = 32,
                                                shuffle = True)

    optimizer = torch.optim.Adam(model.parameters())
    loss_func = nn.CrossEntropyLoss()
    report_loss = 0
    step = 0
    for in_data, label in data_loader_train:
        batch_size = len(in_data)
        optimizer.zero_grad()
        step += 1
        out = model(in_data.view(batch_size, 1, 784))
        loss = loss_func(out, label)
        loss.backward()
        optimizer.step()
        report_loss += loss.item()
        if step % 10 == 0:
            print("report_loss is : " + str(report_loss))
            writer.add_scalar("train/loss", report_loss, step)
            report_loss = 0
        
        
    torch.save(model.state_dict(), "./mnist_model.pkl")
    right_num = 0
    for in_data, label in data_loader_test:
        batch_size = len(in_data)
        out = model(in_data.view(batch_size, 1, 784))
        pred = out.argmax(dim=-1)
        for i, each_pred in enumerate(pred):
            if each_pred == label[i]:
                right_num += 1
        

        # in_data = in_data.view((28, 28))
        # plt.imshow(in_data, cmap="gray")
        # plt.show()
        # print("label is : " + str(out.argmax(dim=-1)))
    
    print(right_num / len(data_test))

        